#encoding: utf-8
#author: walker
# date: 2019-05-06
#summary: 测试gzip压缩/解压文件

import os
import sys
import time
import toml
import gzip


StartTime = time.time()
SrcRoot = r''
DstRoot = r''

def ReadConfig():
    r""" 读取配置文件 """
    global SrcRoot, DstRoot
    cfgFile = 'config.toml'
    if not os.path.exists(cfgFile):
        input(cfgFile + ' not found')
        sys.exit(-1)
    with open(cfgFile, mode='rb') as f:
        content = f.read()
    if content.startswith(b'\xef\xbb\xbf'):     # 去掉 utf8 bom 头
        content = content[3:]
    dic = toml.loads(content.decode('utf8'))

    SrcRoot = dic['SrcRoot'].strip()
    if not os.path.exists(SrcRoot):
        print('Error: not exists %s' % SrcRoot)
        sys.exit(-1)
    print('SrcRoot: %s' % SrcRoot)

    DstRoot = dic['DstRoot'].strip()
    if not os.path.exists(DstRoot):
        print('Error: not exists %s' % DstRoot)
        sys.exit(-1)
    print('DstRoot: %s' % DstRoot)

    print('Read config.toml successed!')


class GZipTool:
    def __init__(self, bufSize):
        self.bufSize = bufSize
        self.fin = None
        self.fout = None

    def compress(self, src, dst):
        self.fin = open(src, 'rb')
        self.fout = gzip.open(dst, 'wb')

        self.__in2out()

    def decompress(self, gzFile, dst):
        self.fin = gzip.open(gzFile, 'rb')
        self.fout = open(dst, 'wb')

        self.__in2out()

    def __in2out(self,):
        while True:
            buf = self.fin.read(self.bufSize)
            if len(buf) < 1:
                break
            self.fout.write(buf)

        self.fin.close()
        self.fout.close()

def Main():
    ReadConfig()
    processed = 0   # 已处理文件个数
    srcSumSize = 0     # 压缩前文件总大小
    dstSumSize = 0     # 压缩后文件总大小
    gzTool = GZipTool(2*1024*1024)
    for parent, dirnames, filenames in os.walk(SrcRoot):
        for filename in filenames:
            if not filename.endswith('.big_json'):
                continue
            srcFile = os.path.join(SrcRoot, filename)
            dstFile = os.path.join(DstRoot, filename + '.gz')
            srcFileSize = os.path.getsize(srcFile)
            srcSumSize += srcFileSize

            print('%s(%.2fMB)\n --->>> \n%s' % (srcFile, srcFileSize/1024/1024, dstFile))
            gzTool.compress(srcFile, dstFile)
            dstFileSize = os.path.getsize(dstFile)
            dstSumSize += dstFileSize
            print('%d B, %.2f MB/s' % (srcSumSize, srcSumSize/1024/1024/(time.time()-StartTime)))
            
            processed += 1
            print('time cost: %.2f s, processed: %d' % (time.time()-StartTime, processed))
            print('')
    
    print('srcSumSize: %d B, %.2f MB, %.2f GB' % (srcSumSize, srcSumSize/1024/1024, srcSumSize/1024/1024/1024))
    print('dstSumSize: %d B, %.2f MB, %.2f GB' % (dstSumSize, dstSumSize/1024/1024, dstSumSize/1024/1024/1024))

if __name__ == '__main__':
    Main()
    print('Time total: %.2f s' % (time.time()-StartTime))
    print(time.strftime('%Y-%m-%d %H:%M:%S', time.localtime()))